"""
Case Type   : 向量数据库可靠性  
Case Name   : ivfflat索引在重建过程注入备机switchover故障
Create At   : 2025/9
Owner       : l1azzzy
Description :
    1、建表
    2、插入大量数据
    3、创建索引
    4、步骤3执行过程中,备机进行switchover
    5、新主重新进行索引创建
    6、主备使用索引查询
    7、主机重新构建索引
    8、主备均使用索引查询
    9、还原主备
    10、清理表
Expect      :
    1、建表成功
    2、插入数据成功
    3、创建索引成功
    4、备机进行switchover成功
    5、新主重新进行索引创建成功
    6、主备使用索引查询成功,主备查询结果相同
    7、主机重新构建索引成功
    8、主备均使用索引查询成功,查询结果相同
    9、还原主备成功
    10、清理表成功
History     :
"""

import os
import re
import time
import threading
import unittest

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

class DatavecIndexCase0032(unittest.TestCase):

    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)} start----')
        self.user_node = [CommonSH('PrimaryDbUser'), CommonSH('Standby1DbUser'), CommonSH('Standby2DbUser')]
        self.pri_user = Node(node='PrimaryDbUser')
        self.standby_user = Node(node='Standby1DbUser')
        self.com = Common('PrimaryDbUser')
        self.com_standby = Common('Standby1DbUser')
        self.constant = Constant()
        self.table_name = 't_datavec_index_case0031'
        self.index_name = 'idx_datavec_index_case0031'
        self.index_created = False
        self.original_primary = 'PrimaryDbUser'
        self.log.info(f'----{os.path.basename(__file__)} end----')

    def trigger_switchover(self):
        """在后台触发备机switchover"""
        time.sleep(10)
        text = '-----step4:步骤3执行过程中,备机进行switchover  expect:备机进行switchover成功-----'
        self.log.info(text)
        switchover_cmd = f'gs_ctl switchover -D {macro.DB_INSTANCE_PATH}'
        res = self.standby_user.sh(switchover_cmd)
        self.log.info(res)
        if 'SIGTERM' in res or 'switchover' in res.lower():
            self.log.info("成功触发备机switchover")
        time.sleep(30)

    def test_datavec_index(self):
        text = '-----step1:建表  expect:建表成功-----'
        self.log.info(text)
        sql = f'create table {self.table_name} (id bigserial primary key, val vector(160));'
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, res)

        text = '-----step2:插入数据  expect:插入数据成功-----'
        self.log.info(text)
        sql = f'''
                    do $$
                    declare
                    vector_dim int := 160;
                    vector_value float[];
                    record_num int := 20000;
                    i int;
                    j int;
                    begin
                    for j in 1..record_num loop
                    for i in 1..vector_dim loop
                    vector_value[i] := ROUND(RANDOM() * 100, 2);
                    end loop;
                    insert into {self.table_name} (val) values (vector_value);
                    end loop;
                    end $$;
                    '''
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, res)

        text = '-----step3:创建索引  expect:创建索引成功-----'
        self.log.info(text)
        switchover_thread = threading.Thread(target=self.trigger_switchover)
        switchover_thread.daemon = True
        switchover_thread.start()

        sql = f'''
        alter table {self.table_name} set(parallel_workers=30);
        set maintenance_work_mem ='1GB';
        create index {self.index_name} on {self.table_name} using ivfflat (val vector_l2_ops) with(m=100,ef_construction=200);
        '''
        try:
            res = self.com.execut_db_sql(sql)
            self.log.info(res)
            if self.constant.INDEX_CREATE_SUCCESS in res:
                self.index_created = True
                self.log.info("索引创建在switchover前已完成")
        except Exception as e:
            self.log.info(f"索引创建过程执行异常: {e}")
            if "connection" in str(e).lower() or "closed" in str(e).lower():
                self.log.info("索引创建因switchover而中断")
        switchover_thread.join(timeout=60)

        text = '-----step5:新主重新进行索引创建  expect:新主重新进行索引创建成功-----'
        self.log.info(text)
        time.sleep(30)
        max_retries = 10
        new_primary_com = None
        for i in range(max_retries):
            try:
                check_sql = 'select 1;'
                res = self.com.execut_db_sql(check_sql)
                if '1' in res:
                    new_primary_com = self.com
                    self.log.info("原主机仍是主节点")
                    break
            except Exception:
                try:
                    check_sql = 'select 1;'
                    res = self.com_standby.execut_db_sql(check_sql)
                    if '1' in res:
                        new_primary_com = self.com_standby
                        self.log.info("原备机已成为新主节点")
                        break
                except Exception:
                    self.log.info(f"寻找新主节点，重试 {i+1}/{max_retries}")
                    time.sleep(10)
        else:
            self.fail("无法找到新的主节点")
        if not self.index_created:
            sql = f'create index {self.index_name} on {self.table_name} using ivfflat (val vector_l2_ops) with(m=100,ef_construction=200);'
            res = new_primary_com.execut_db_sql(sql)
            self.log.info(res)
            self.assertIn(self.constant.INDEX_CREATE_SUCCESS, res)
            self.index_created = True
        else:
            self.log.info("索引已存在，无需重新创建")

        text = '-----step6:主备使用索引查询  expect:主备使用索引查询成功,主备查询结果相同-----'
        self.log.info(text)
        query_sql = f'''
        set enable_indexscan=on;
        set enable_seqscan=off;
        select id,val<->'[33.91,98.59,85.29,15.32,4.9,
        78.51,15.83,1.63,39.63,52,40.46,74.03,16.95,95.42,
        15.85,56.45,28.41,69.76,55.21,68.63,53.24,75.3,5.87,
        42.23,26.98,81.97,89.8,58.99,0.04,85.03,73.05,1.27,
        59.57,79.81,87.65,81.77,34,3.57,6.12,34.42,43.21,93.49,
        81.3,64.9,87.38,0.36,2.03,2.05,54.29,43.58,62.91,76.71,
        47.7,14.25,6.93,9.56,8.46,23.97,24.78,84.65,86.24,55.93,
        2.43,50,73.33,23.49,15.29,44.7,94.62,56.72,8.84,36.69,
        50.54,69.76,62.04,41.83,84.16,64.97,53.27,22.08,22.78,
        83.52,85.29,54.41,68.86,47.64,42.35,65.58,59.54,64.85,61.89,
        21.02,9.73,89.1,87.04,82.73,15.58,93.47,3.27,46.71,45.06,91.6,
        49.84,80.88,35.7,14.65,78.72,97.65,70,73.04,23.79,62.36,0.7,99.44,
        50.06,98.07,56.67,25.9,6.77,55.16,10.47,13.38,82.87,61.39,25.64,
        48.81,4.9,13.26,2.41,1.46,65.29,58.76,89.64,86.72,84.13,13.96,
        45.98,98.53,95.92,52.29,90.25,19.42,90.89,68.13,37.44,4.48,
        99.26,98.01,40.22,18.98,78.94,50.45,42.93,21.05,89.19,53.91,
        65.14,56.68,12.69,51.74]' from {self.table_name} 
        order by val<->'[33.91,98.59,85.29,15.32,4.9,78.51,15.83,1.63,
        39.63,52,40.46,74.03,16.95,95.42,15.85,56.45,28.41,69.76,55.21,
        68.63,53.24,75.3,5.87,42.23,26.98,81.97,89.8,58.99,0.04,85.03,
        73.05,1.27,59.57,79.81,87.65,81.77,34,3.57,6.12,34.42,43.21,
        93.49,81.3,64.9,87.38,0.36,2.03,2.05,54.29,43.58,62.91,76.71,
        47.7,14.25,6.93,9.56,8.46,23.97,24.78,84.65,86.24,55.93,2.43,50,
        73.33,23.49,15.29,44.7,94.62,56.72,8.84,36.69,50.54,69.76,62.04,
        41.83,84.16,64.97,53.27,22.08,22.78,83.52,85.29,54.41,68.86,47.64,
        42.35,65.58,59.54,64.85,61.89,21.02,9.73,89.1,87.04,82.73,15.58,
        93.47,3.27,46.71,45.06,91.6,49.84,80.88,35.7,14.65,78.72,
        97.65,70,73.04,23.79,62.36,0.7,99.44,50.06,98.07,56.67,25.9,
        6.77,55.16,10.47,13.38,82.87,61.39,25.64,48.81,4.9,13.26,2.41,
        1.46,65.29,58.76,89.64,86.72,84.13,13.96,45.98,98.53,95.92,52.29,
        0.25,19.42,90.89,68.13,37.44,4.48,99.26,98.01,40.22,18.98,78.94,
        50.45,42.93,21.05,89.19,53.91,65.14,56.68,12.69,51.74]' limit 10;
        '''
        res_new_primary = new_primary_com.execut_db_sql(query_sql)
        self.log.info(f"新主查询结果: {res_new_primary}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_new_primary)
        try:
            res_new_standby = self.com.execut_db_sql(query_sql) if new_primary_com != self.com else self.com_standby.execut_db_sql(query_sql)
            self.log.info(f"新备查询结果: {res_new_standby}")
            self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_new_standby)
        except Exception as e:
            self.log.info(f"新备查询异常（可能仍在恢复中）: {e}")

        text = '-----step7:主机重新构建索引  expect:主机重新构建索引成功-----'
        self.log.info(text)
        sql = f'reindex index {self.index_name};'
        res = new_primary_com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.INDEX_CREATE_SUCCESS, res)

        text = '-----step8:主备均使用索引查询  expect:主备均使用索引查询成功,查询结果相同-----'
        self.log.info(text)
        res_new_primary_after = new_primary_com.execut_db_sql(query_sql)
        self.log.info(f"重建索引后新主查询结果: {res_new_primary_after}")
        self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_new_primary_after)
        try:
            res_new_standby_after = self.com.execut_db_sql(query_sql) if new_primary_com != self.com else self.com_standby.execut_db_sql(query_sql)
            self.log.info(f"重建索引后新备查询结果: {res_new_standby_after}")
            self.assertIn(self.constant.SELECT_SUCCESS_MSG, res_new_standby_after)
        except Exception as e:
            self.log.info(f"新备查询异常: {e}")

    def tearDown(self):
        text = '-----step9:还原主备  expect:还原主备成功-----'
        self.log.info(text)
        try:
            switchover_cmd = f'gs_ctl switchover -D {macro.DB_INSTANCE_PATH}'
            res = self.standby_user.sh(switchover_cmd)
            self.log.info(res)
            self.assertTrue('SIGTERM' in res or 'switchover' in res.lower(), "还原主备失败")
            time.sleep(30)
        except Exception as e:
            self.log.info(f"还原主备过程中出现异常: {e}")

        text = '-----step10:清理表  expect:清理表成功-----'
        self.log.info(text)
        sql = f'drop table if exists {self.table_name};'
        res = self.com.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_TABLE_SUCCESS_MSG, res)
        self.log.info(f'----{os.path.basename(__file__)} end----')

    def extract_query_results(self, query_result):
        """从查询结果中提取id和距离值用于比较"""
        lines = query_result.split('\n')
        results = []
        for line in lines:
            if '|' in line and not line.startswith('---') and not line.startswith(' id ') and line.strip():
                parts = line.split('|')
                if len(parts) >= 2:
                    try:
                        id_val = parts[0].strip()
                        distance = parts[1].strip()
                        if id_val.isdigit():
                            results.append((int(id_val), float(distance)))
                    except (ValueError, IndexError):
                        continue
        return sorted(results)